package com.xiashitech.agent.instrumentation.template.inst;

import com.xiashitech.agent.instrumentation.base.InstPreparedArgs;
import com.xiashitech.agent.instrumentation.base.inst.CommonInstrumentation;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSLettuceCallDepth;
import com.xiashitech.agent.instrumentation.template.TemplateRequestData;
import com.xiashitech.agent.instrumentation.template.TemplateSingletons;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
import static net.bytebuddy.matcher.ElementMatchers.*;

/**
 * 插桩织入示例代码
 */
public class TemplateInstrumentation extends CommonInstrumentation {

    @Override
    public ElementMatcher<ClassLoader> classLoaderOptimization() {
        //如果织入Class的ClassLoader不是UrlClassLoader（比如bootstrap/platform），该方法无法匹配
        return hasClassesNamed("template.Template");
    }

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return implementsInterface(named("template.Template"));
    }

    @Override
    public void transform(TypeTransformer typeTransformer) {
        transformMethod(isMethod()
                        .and(named("get"))
                        .and(takesArguments(1))
                        .and(isPublic()),
                "GetAdvice", typeTransformer);
    }

    /**
     * Advice中仅做流程调用以及必要的REQUEST/RESPONSE对象构建操作，业务代码统一放在processor中执行
     */
    public static class GetAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.This Object obj,
                @Advice.Argument(0) String param1,
                @Advice.Local("preparedArgs") InstPreparedArgs preparedArgs,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("requestData") TemplateRequestData requestData) {

            //onEnter时的检查和准备
            preparedArgs = InstPreparedArgs.prepareOnEnter("Template.get", false, XSLettuceCallDepth.class);
            if(!InstPreparedArgs.checkOnEnter(preparedArgs))
                return;

            requestData = new TemplateRequestData(); //构建请求数据
            context = preparedArgs.getContext(); //从preparedArgs中获取context(已经和TraceContextVars同步过)
            if(!TemplateSingletons.instrumenter().shouldStart(context, requestData))
                return;

            //先执行start，再创建scope
            context = TemplateSingletons.instrumenter().start(context, requestData);
            scope = context.makeCurrent();
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Return Object rtn,
                @Advice.Local("preparedArgs") InstPreparedArgs preparedArgs,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("requestData") TemplateRequestData requestData) {

            //onExit时的检查和准备
            if(!InstPreparedArgs.checkOnExit(preparedArgs))
                return;
            if(scope == null)
                return;

            //先关闭scope，再执行end
            scope.close();
            TemplateSingletons.instrumenter().end(context, requestData, rtn, null);
        }
    }

}
